import uvicorn

from mcp.server.fastmcp import FastMCP
from mcp.server.sse import SseServerTransport

from starlette.applications import Starlette
from starlette.requests import Request
from starlette.routing import Mount, Route


mcp = FastMCP("mcp server demo...")

@mcp.tool()
def weather(address: str) -> str:
    """ 查询某个地区的天气情况

    Args:
        address: 地区名称
    """
    if address == '深圳':
        return '晴朗'
    elif address == '广州':
        return '暴雨'
    else:
        return '多云'


@mcp.tool()
def transportation(address: str):
    """ 查询某个地区的交通情况

    Args:
        address: 地区名称
    """
    if address == '深圳':
        return '道路顺畅'
    elif address == '广州':
        return '道路拥堵'
    else:
        return '未知'

def get_starlette_app(mcp_server, sse):
    async def handle_sse(request: Request) -> None:
        async with sse.connect_sse(
                request.scope,
                request.receive,
                request._send,
        ) as (read_stream, write_stream):
            await mcp_server.run(
                read_stream,
                write_stream,
                mcp_server.create_initialization_options(),
            )

    return Starlette(
        debug=True,
        routes=[
            Route("/sse", endpoint=handle_sse),
            Mount("/messages/", app=sse.handle_post_message),
        ]
    )


class McpServer:
    @staticmethod
    def start():
        mcp_server, sse = mcp._mcp_server, SseServerTransport("/messages/")
        uvicorn.run(get_starlette_app(mcp_server, sse), host="0.0.0.0", port=8000)